import numpy as np
import matplotlib.pyplot as plt
import os
import re

PATH = '../../../../large_data/DL1/cifar-10-batches-bin/'
PIC_SIDE = 32
PIC_CH = 3
LABEL_LEN = 1
RECORED_SIZE = LABEL_LEN + PIC_SIDE ** 2 * PIC_CH


def read_bin_data(m):
    with open(PATH + 'data_batch_1.bin', 'br') as f:
        buf = f.read(m * RECORED_SIZE)
        data = np.frombuffer(buf, dtype=np.uint8).reshape([-1, RECORED_SIZE])
        labels, images = np.hsplit(data, [LABEL_LEN])
        images = images.reshape([-1, PIC_CH, PIC_SIDE, PIC_SIDE])
        images = images.transpose([0, 2, 3, 1])
        return labels, images


with open(PATH + 'batches.meta.txt', 'r') as f:
    labels = f.readlines()
    labels = [re.sub(r'[\r\n]', '', s) for s in labels]
    idx2labels = []
    for s in labels:
        if len(s) == 0:
            continue
        idx2labels.append(s)

if '__main__' == __name__:
    import matplotlib.pyplot as plt

    plt.figure(figsize=[16, 8])
    spr = 4  # subplot row
    spc = 8  # subplot column
    spn = 0

    y, x = read_bin_data(spr * spc)
    for i, pic in enumerate(x):
        spn += 1
        if spn > spr * spc:
            break
        plt.subplot(spr, spc, spn)
        plt.axis('off')
        idx = y[i, 0]
        plt.title(f'"{idx2labels[idx]}"')  # ATTENTION y is 10x1, y[i] is vector, y[i, 0] is scalar, this diff especially important when labels is a pandas Series
        plt.imshow(pic)